import numpy as np
import cv2
import copy


def offset_to_lengths(lod):
    offset = lod[0]
    lengths = [offset[i + 1] - offset[i] for i in range(len(offset) - 1)]
    return [lengths]


def get_extra_info(im, arch, shape, scale):
    print(shape)
    info = []
    if 'YOLO' in arch:
        im_size = np.array([shape[:2]]).astype('int32')
        info.append(im_size)
    return info


class Resize(object):
    def __init__(self, target_size, max_size=0, interp=cv2.INTER_LINEAR):
        super(Resize, self).__init__()
        self.target_size = target_size
        self.max_size = max_size
        self.interp = interp

    def __call__(self, im, arch):
        origin_shape = im.shape[:2]
        im_c = im.shape[2]
        scale_set = {'RCNN', 'RetinaNet'}
        if self.max_size != 0 and arch in scale_set:
            im_size_min = np.min(origin_shape[0:2])
            im_size_max = np.max(origin_shape[0:2])
            im_scale = float(self.target_size) / float(im_size_min)
            if np.round(im_scale * im_size_max) > self.max_size:
                im_scale = float(self.max_size) / float(im_size_max)
            im_scale_x = im_scale
            im_scale_y = im_scale
            resize_w = int(im_scale_x * float(origin_shape[1]))
            resize_h = int(im_scale_y * float(origin_shape[0]))
        else:
            im_scale_x = float(self.target_size) / float(origin_shape[1])
            im_scale_y = float(self.target_size) / float(origin_shape[0])
        im = cv2.resize(
            im,
            None,
            None,
            fx=im_scale_x,
            fy=im_scale_y,
            interpolation=self.interp)
        # padding im
        if self.max_size != 0 and arch in scale_set:
            padding_im = np.zeros(
                (self.max_size, self.max_size, im_c), dtype=np.float32)
            im_h, im_w = im.shape[:2]
            padding_im[:im_h, :im_w, :] = im
            im = padding_im
        return im, im_scale_x


class Normalize(object):
    def __init__(self, mean, std, is_scale=True):
        super(Normalize, self).__init__()
        self.mean = mean
        self.std = std
        self.is_scale = is_scale

    def __call__(self, im):
        im = im.astype(np.float32, copy=False)
        if self.is_scale:
            im = im / 255.0
        im -= self.mean
        im /= self.std
        return im


class Permute(object):
    def __init__(self, to_bgr=False):
        self.to_bgr = to_bgr

    def __call__(self, im):
        im = im.transpose((2, 0, 1)).copy()
        if self.to_bgr:
            im = im[[2, 1, 0], :, :]
        return im


class PadStride(object):
    def __init__(self, stride=0):
        # assert stride >= 0, "Unsupported stride: {}, the stride in PadStride must be greater or equal to 0".format(
        #     stride)
        self.coarsest_stride = stride

    def __call__(self, im):
        coarsest_stride = self.coarsest_stride
        if coarsest_stride == 0:
            return im
        im_c, im_h, im_w = im.shape
        pad_h = int(np.ceil(float(im_h) / coarsest_stride) * coarsest_stride)
        pad_w = int(np.ceil(float(im_w) / coarsest_stride) * coarsest_stride)
        padding_im = np.zeros((im_c, pad_h, pad_w), dtype=np.float32)
        padding_im[:, :im_h, :im_w] = im
        return padding_im


def Preprocess(img_path, arch, config):
    img = img_path
    # img = DecodeImage(img_path)
    orig_shape = img.shape
    scale = 1.
    data = []
    data_config = copy.deepcopy(config)
    for data_aug_conf in data_config:
        obj = data_aug_conf.pop('type')
        preprocess = eval(obj)(**data_aug_conf)
        if obj == 'Resize':
            img, scale = preprocess(img, arch)
        else:
            img = preprocess(img)

    img = img[np.newaxis, :]  # N, C, H, W
    data.append(img)
    extra_info = get_extra_info(img, arch, orig_shape, scale)
    data += extra_info
    return data


def coco17_category_info(with_background=True):
    """
    Get class id to category id map and category id
    to category name map of COCO2017 dataset

    Args:
        with_background (bool, default True):
            whether load background as class 0.
    """
    clsid2catid = {
        1: 1,
        2: 2,
        3: 3,
        4: 4,
        5: 5,
        6: 6,
        7: 7,
        8: 8,
        9: 9,
        10: 10,
        11: 11,
        12: 13,
        13: 14,
        14: 15,
        15: 16,
        16: 17,
        17: 18,
        18: 19,
        19: 20,
        20: 21,
        21: 22,
        22: 23,
        23: 24,
        24: 25,
        25: 27,
        26: 28,
        27: 31,
        28: 32,
        29: 33,
        30: 34,
        31: 35,
        32: 36,
        33: 37,
        34: 38,
        35: 39,
        36: 40,
        37: 41,
        38: 42,
        39: 43,
        40: 44,
        41: 46,
        42: 47,
        43: 48,
        44: 49,
        45: 50,
        46: 51,
        47: 52,
        48: 53,
        49: 54,
        50: 55,
        51: 56,
        52: 57,
        53: 58,
        54: 59,
        55: 60,
        56: 61,
        57: 62,
        58: 63,
        59: 64,
        60: 65,
        61: 67,
        62: 70,
        63: 72,
        64: 73,
        65: 74,
        66: 75,
        67: 76,
        68: 77,
        69: 78,
        70: 79,
        71: 80,
        72: 81,
        73: 82,
        74: 84,
        75: 85,
        76: 86,
        77: 87,
        78: 88,
        79: 89,
        80: 90
    }

    catid2name = {
        0: 'background',
        1: 'person',
        2: 'bicycle',
        3: 'car',
        4: 'motorcycle',
        5: 'airplane',
        6: 'bus',
        7: 'train',
        8: 'truck',
        9: 'boat',
        10: 'traffic light',
        11: 'fire hydrant',
        13: 'stop sign',
        14: 'parking meter',
        15: 'bench',
        16: 'bird',
        17: 'cat',
        18: 'dog',
        19: 'horse',
        20: 'sheep',
        21: 'cow',
        22: 'elephant',
        23: 'bear',
        24: 'zebra',
        25: 'giraffe',
        27: 'backpack',
        28: 'umbrella',
        31: 'handbag',
        32: 'tie',
        33: 'suitcase',
        34: 'frisbee',
        35: 'skis',
        36: 'snowboard',
        37: 'sports ball',
        38: 'kite',
        39: 'baseball bat',
        40: 'baseball glove',
        41: 'skateboard',
        42: 'surfboard',
        43: 'tennis racket',
        44: 'bottle',
        46: 'wine glass',
        47: 'cup',
        48: 'fork',
        49: 'knife',
        50: 'spoon',
        51: 'bowl',
        52: 'banana',
        53: 'apple',
        54: 'sandwich',
        55: 'orange',
        56: 'broccoli',
        57: 'carrot',
        58: 'hot dog',
        59: 'pizza',
        60: 'donut',
        61: 'cake',
        62: 'chair',
        63: 'couch',
        64: 'potted plant',
        65: 'bed',
        67: 'dining table',
        70: 'toilet',
        72: 'tv',
        73: 'laptop',
        74: 'mouse',
        75: 'remote',
        76: 'keyboard',
        77: 'cell phone',
        78: 'microwave',
        79: 'oven',
        80: 'toaster',
        81: 'sink',
        82: 'refrigerator',
        84: 'book',
        85: 'clock',
        86: 'vase',
        87: 'scissors',
        88: 'teddy bear',
        89: 'hair drier',
        90: 'toothbrush'
    }

    if not with_background:
        clsid2catid = {k - 1: v for k, v in clsid2catid.items()}

    return clsid2catid, catid2name


def oid_category_info():
    """
    Get class id to category id map and category id
    to category name map of oid dataset

    Args:
        with_background (bool, default True):
            whether load background as class 0.
    """
    catid2name = {
        0: 'Building',
        1: 'Car',
        2: 'Taxi',
        3: 'Human eye',
        4: 'Shirt',
        5: 'Human mouth',
        6: 'Human ear',
        7: 'Human hair',
        8: 'Human head',
        9: 'Man',
        10: 'Human face',
        11: 'Human arm',
        12: 'Glasses',
        13: 'Human hand',
        14: 'Bronze sculpture',
        15: 'Person',
        16: 'Shotgun',
        17: 'Helmet',
        18: 'Curtain',
        19: 'Table',
        20: 'Backpack',
        21: 'Vehicle',
        22: 'Carnivore',
        23: 'Tree',
        24: 'Miniskirt',
        25: 'Woman',
        26: 'Footwear',
        27: 'Coat',
        28: 'Scarf',
        29: 'Jacket',
        30: 'Human leg',
        31: 'Jeans',
        32: 'Lantern',
        33: 'Window',
        34: 'Handbag',
        35: 'Light bulb',
        36: 'Mule',
        37: 'Traffic light',
        38: 'Bus',
        39: 'Office building',
        40: 'Wine',
        41: 'Wine glass',
        42: 'Home appliance',
        43: 'Laptop',
        44: 'Computer keyboard',
        45: 'Printer',
        46: 'Chair',
        47: 'Cabinetry',
        48: 'Desk',
        49: 'Computer monitor',
        50: 'Office supplies',
        51: 'House',
        52: 'Chest of drawers',
        53: 'Cupboard',
        54: 'Drawer',
        55: 'Shelf',
        56: 'Coffee cup',
        57: 'Coffee table',
        58: 'Countertop',
        59: 'Microwave oven',
        60: 'Sunglasses',
        61: 'Shorts',
        62: 'Truck',
        63: 'Wheel',
        64: 'Van',
        65: 'Fountain',
        66: 'Sculpture',
        67: 'Human foot',
        68: 'Bookcase',
        69: 'Book',
        70: 'Dinosaur',
        71: 'Parrot',
        72: 'Bicycle',
        73: 'Camera',
        74: 'Rifle',
        75: 'Watermelon',
        76: 'Guitar',
        77: 'Violin',
        78: 'Dress',
        79: 'Suit',
        80: 'Girl',
        81: 'Flower',
        82: 'Train',
        83: 'Boy',
        84: 'Mug',
        85: 'Drinking straw',
        86: 'Beer',
        87: 'Cocktail',
        88: 'Mobile phone',
        89: 'Swimming pool',
        90: 'Dog',
        91: 'Trousers',
        92: 'Furniture',
        93: 'Tire',
        94: 'Human nose',
        95: 'Castle',
        96: 'Fish',
        97: 'Marine mammal',
        98: 'Swimwear',
        99: 'Brassiere',
        100: 'Cattle',
        101: 'Goat',
        102: 'Sheep',
        103: 'Adhesive tape',
        104: 'Hat',
        105: 'Crutch',
        106: 'Squirrel',
        107: 'Street light',
        108: 'Sparrow',
        109: 'Gas stove',
        110: 'Tableware',
        111: 'Kitchen appliance',
        112: 'Stretcher',
        113: 'Squash',
        114: 'Cutting board',
        115: 'Fruit',
        116: 'Potato',
        117: 'Carrot',
        118: 'Scoreboard',
        119: 'Billboard',
        120: 'Sports uniform',
        121: 'Couch',
        122: 'Door',
        123: 'Horse',
        124: 'Alpaca',
        125: 'Digital clock',
        126: 'Headphones',
        127: 'Helicopter',
        128: 'Waste container',
        129: 'Human beard',
        130: 'Stool',
        131: 'Penguin',
        132: 'Invertebrate',
        133: 'Frog',
        134: 'Sombrero',
        135: 'Doll',
        136: 'Baseball bat',
        137: 'Bottle',
        138: 'Plate',
        139: 'Serving tray',
        140: 'Spoon',
        141: 'Cake',
        142: 'Refrigerator',
        143: 'Ice cream',
        144: 'Mirror',
        145: 'Tent',
        146: 'Falcon',
        147: 'Bowl',
        148: 'Cannon',
        149: 'Insect',
        150: 'Orange',
        151: 'Vegetable',
        152: 'Grapefruit',
        153: 'Surfboard',
        154: 'Whiteboard',
        155: 'Sock',
        156: 'Sandal',
        157: 'Baseball glove',
        158: 'Vehicle registration plate',
        159: 'Toy',
        160: 'Dragonfly',
        161: 'Boot',
        162: 'Saucer',
        163: 'Fork',
        164: 'Drum',
        165: 'Microphone',
        166: 'Fedora',
        167: 'Sun hat',
        168: 'Flag',
        169: 'Skirt',
        170: 'Bicycle wheel',
        171: 'Bicycle helmet',
        172: 'Snail',
        173: 'Tie',
        174: 'Bench',
        175: 'Umbrella',
        176: 'Tower',
        177: 'Skyscraper',
        178: 'High heels',
        179: 'Christmas tree',
        180: 'Paddle',
        181: 'Canoe',
        182: 'Cookie',
        183: 'Drink',
        184: 'Coffee',
        185: 'Handgun',
        186: 'Tea',
        187: 'Dessert',
        188: 'Strawberry',
        189: 'Bread',
        190: 'Kitchen & dining room table',
        191: 'Sushi',
        192: 'Tin can',
        193: 'Necklace',
        194: 'Earrings',
        195: 'Spider',
        196: 'Traffic sign',
        197: 'Land vehicle',
        198: 'Goggles',
        199: 'Watercraft',
        200: 'Ant',
        201: 'Bee',
        202: 'Turtle',
        203: 'Musical instrument',
        204: 'Cowboy hat',
        205: 'Lamp',
        206: 'Crocodile',
        207: 'Parachute',
        208: 'Marine invertebrates',
        209: 'Organ',
        210: 'Bull',
        211: 'Juice',
        212: 'Airplane',
        213: 'Clock',
        214: 'Barge',
        215: 'Knife',
        216: 'Animal',
        217: 'Balloon',
        218: 'Wheelchair',
        219: 'Tripod',
        220: 'Plastic bag',
        221: 'Football',
        222: 'Tap',
        223: 'Bathtub',
        224: 'Glove',
        225: 'Lily',
        226: 'Tomato',
        227: 'Sink',
        228: 'Houseplant',
        229: 'Motorcycle',
        230: 'Tennis racket',
        231: 'Seat belt',
        232: 'Piano',
        233: 'Cat',
        234: 'Football helmet',
        235: 'Bathroom cabinet',
        236: 'Jellyfish',
        237: 'Barrel',
        238: 'Television',
        239: 'Rose',
        240: 'Bird',
        241: 'Platter',
        242: 'Elephant',
        243: 'Weapon',
        244: 'Bed',
        245: 'Dog bed',
        246: 'Jaguar',
        247: 'Leopard',
        248: 'Cheetah',
        249: 'Sofa bed',
        250: 'Pumpkin',
        251: 'Stairs',
        252: 'Deer',
        253: 'Dolphin',
        254: 'Lighthouse',
        255: 'Personal care',
        256: 'Flowerpot',
        257: 'Ball',
        258: 'Snowboard',
        259: 'Ski',
        260: 'Pig',
        261: 'Pillow',
        262: 'Tiara',
        263: 'Crown',
        264: 'Screwdriver',
        265: 'Poster',
        266: 'Seahorse',
        267: 'Mushroom',
        268: 'Window blind',
        269: 'Box',
        270: 'Tortoise',
        271: 'Ladybug',
        272: 'Bat',
        273: 'Fireplace',
        274: 'Loveseat',
        275: 'Grape',
        276: 'Watch',
        277: 'Skull',
        278: 'Bidet',
        279: 'Toilet',
        280: 'Eagle',
        281: 'Tank',
        282: 'Shellfish',
        283: 'Chicken',
        284: 'Palm tree',
        285: 'Giraffe',
        286: 'Trumpet',
        287: 'Tennis ball',
        288: 'Candy',
        289: 'Convenience store',
        290: 'Sea turtle',
        291: 'Beetle',
        292: 'Seafood',
        293: 'Salad',
        294: 'Teddy bear',
        295: 'Skateboard',
        296: 'Binoculars',
        297: 'Picture frame',
        298: 'Studio couch',
        299: 'Porch',
        300: 'Pizza',
        301: 'Hamburger',
        302: 'Coin',
        303: 'Gondola',
        304: 'Apple',
        305: 'Oven',
        306: 'Nightstand',
        307: 'Luggage and bags',
        308: 'Duck',
        309: 'Swan',
        310: 'Broccoli',
        311: 'Tick',
        312: 'Snake',
        313: 'Wood-burning stove',
        314: 'Boat',
        315: 'Cucumber',
        316: 'Zucchini',
        317: 'Golf ball',
        318: 'Muffin',
        319: 'Candle',
        320: 'Pen',
        321: 'Sandwich',
        322: 'Computer mouse',
        323: 'Tablet computer',
        324: 'Monkey',
        325: 'Fire hydrant',
        326: 'Missile',
        327: 'Briefcase',
        328: 'Vase',
        329: 'Mechanical fan',
        330: 'Ceiling fan',
        331: 'Snowman',
        332: 'Filing cabinet',
        333: 'Fox',
        334: 'Suitcase',
        335: 'Limousine',
        336: 'Rabbit',
        337: 'Aircraft',
        338: 'French fries',
        339: 'Volleyball',
        340: 'Lemon',
        341: 'Ladder',
        342: 'Burrito',
        343: 'Beehive',
        344: 'Honeycomb',
        345: 'Pasta',
        346: 'Sunflower',
        347: 'Lizard',
        348: 'Hamster',
        349: 'Mouse',
        350: 'Butterfly',
        351: 'Goose',
        352: 'Moths and butterflies',
        353: 'Coconut',
        354: 'Woodpecker',
        355: 'Cake stand',
        356: 'Shower',
        357: 'Corded phone',
        358: 'Ruler',
        359: 'Owl',
        360: 'Sword',
        361: 'Washing machine',
        362: 'Picnic basket',
        363: 'Kangaroo',
        364: 'Teapot',
        365: 'Pitcher',
        366: 'Jug',
        367: 'Radish',
        368: 'Dagger',
        369: 'Lion',
        370: 'Hot dog',
        371: 'Croissant',
        372: 'Lavender',
        373: 'Belt',
        374: 'Wrench',
        375: 'Pancake',
        376: 'Maple',
        377: 'Wok',
        378: 'Shrimp',
        379: 'Kite',
        380: 'Door handle',
        381: 'Flashlight',
        382: 'Torch',
        383: 'Whale',
        384: 'Bust',
        385: 'Accordion',
        386: 'Stop sign',
        387: 'Tiger',
        388: 'Chopsticks',
        389: 'Guacamole',
        390: 'Roller skates',
        391: 'Paper towel',
        392: 'Toilet paper',
        393: 'Musical keyboard',
        394: 'Pomegranate',
        395: 'Bell pepper',
        396: 'Egg',
        397: 'Cello',
        398: 'Ostrich',
        399: 'Scissors',
        400: 'Alarm clock',
        401: 'Zebra',
        402: 'Rhinoceros',
        403: 'Crab',
        404: 'Shark',
        405: 'Camel',
        406: 'Oyster',
        407: 'Envelope',
        408: 'Ambulance',
        409: 'Mango',
        410: 'Tart',
        411: 'Flute',
        412: 'Table tennis racket',
        413: 'Telephone',
        414: 'Nail',
        415: 'Sea lion',
        416: 'Cart',
        417: 'Peach',
        418: 'Bow and arrow',
        419: 'Goldfish',
        420: 'Blue jay',
        421: 'Lifejacket',
        422: 'Billiard table',
        423: 'Food processor',
        424: 'Mixer',
        425: 'Towel',
        426: 'Trombone',
        427: 'Taco',
        428: 'Cabbage',
        429: 'Antelope',
        430: 'Racket',
        431: 'Rocket',
        432: 'Banana',
        433: 'Reptile',
        434: 'Frying pan',
        435: 'Bear',
        436: 'Starfish',
        437: 'Segway',
        438: 'Kitchen knife',
        439: 'Coffeemaker',
        440: 'Wall clock',
        441: 'Saxophone',
        442: 'Toaster',
        443: 'Raven',
        444: 'Stationary bicycle',
        445: 'Popcorn',
        446: 'Light switch',
        447: 'Centipede',
        448: 'Caterpillar',
        449: 'Golf cart',
        450: 'Swim cap',
        451: 'Harbor seal',
        452: 'Sewing machine',
        453: 'Bagel',
        454: 'Doughnut',
        455: 'Brown bear',
        456: 'Harpsichord',
        457: 'Jet ski',
        458: 'Infant bed',
        459: 'Polar bear',
        460: 'Punching bag',
        461: 'Lobster',
        462: 'Canary',
        463: 'Turkey',
        464: 'Raccoon',
        465: 'Plumbing fixture',
        466: 'Rugby ball',
        467: 'Common fig',
        468: 'Kettle',
        469: 'Pear',
        470: 'Artichoke',
        471: 'Snowplow',
        472: 'Waffle',
        473: 'Otter',
        474: 'Porcupine',
        475: 'Pineapple',
        476: 'Asparagus',
        477: 'Cricket ball',
        478: 'Slow cooker',
        479: 'Measuring cup',
        480: 'Beaker',
        481: 'Harp',
        482: 'Ring binder',
        483: 'Horn',
        484: 'Willow',
        485: 'Training bench',
        486: 'Snowmobile',
        487: 'Submarine sandwich',
        488: 'Power plugs and sockets',
        489: 'Oboe',
        490: 'Blender',
        491: 'Spatula',
        492: 'Dice',
        493: 'Winter melon',
        494: 'Treadmill',
        495: 'Pretzel',
        496: 'Salt and pepper shakers',
        497: 'Lynx',
        498: 'Dumbbell',
        499: 'Pressure cooker'
    }

    return catid2name


def clip_bbox(bbox):
    xmin = max(min(bbox[0], 1.), 0.)
    ymin = max(min(bbox[1], 1.), 0.)
    xmax = max(min(bbox[2], 1.), 0.)
    ymax = max(min(bbox[3], 1.), 0.)
    return xmin, ymin, xmax, ymax


def bbox2out(results, clsid2catid, is_bbox_normalized=False):
    """
    Args:
        results: request a dict, should include: `bbox`, `im_id`,
                 if is_bbox_normalized=True, also need `im_shape`.
        clsid2catid: class id to category id map of COCO2017 dataset.
        is_bbox_normalized: whether or not bbox is normalized.
    """
    xywh_res = []
    for t in results:
        bboxes = t['bbox'][0]
        lengths = t['bbox'][1][0]
        im_ids = np.array(t['im_id'][0]).flatten()
        if bboxes.shape == (1, 1) or bboxes is None:
            continue

        k = 0
        for i in range(len(lengths)):
            num = lengths[i]
            im_id = int(im_ids[i])
            for j in range(num):
                dt = bboxes[k]
                clsid, score, xmin, ymin, xmax, ymax = dt.tolist()
                catid = (clsid2catid[int(clsid)])

                if is_bbox_normalized:
                    xmin, ymin, xmax, ymax = \
                            clip_bbox([xmin, ymin, xmax, ymax])
                    w = xmax - xmin
                    h = ymax - ymin
                    im_shape = t['im_shape'][0][i].tolist()
                    im_height, im_width = int(im_shape[0]), int(im_shape[1])
                    xmin *= im_width
                    ymin *= im_height
                    w *= im_width
                    h *= im_height
                else:
                    w = xmax - xmin + 1
                    h = ymax - ymin + 1

                bbox = [xmin, ymin, w, h]
                coco_res = {
                    'image_id': im_id,
                    'category_id': catid,
                    'bbox': bbox,
                    'score': score
                }
                xywh_res.append(coco_res)
                k += 1
    return xywh_res
